import streamlit as st
import sqlite3
import hashlib
import os

db_dir = "db"
db_file = "users.db"
db_path = os.path.join(db_dir, db_file)
if not os.path.exists(db_dir):
    os.makedirs(db_dir)
conn = sqlite3.connect(db_path,check_same_thread=False)
c = conn.cursor()

def create_users_table():
    c.execute('CREATE TABLE IF NOT EXISTS users(username TEXT, password TEXT)')

def add_user(username, password):
    """向数据库添加新用户"""
    c.execute('INSERT INTO users(username, password) VALUES (?, ?)', (username, hash_password(password)))
    conn.commit()

def login_user(username, password):
    """从数据库验证用户登录信息"""
    c.execute('SELECT * FROM users WHERE username = ? AND password = ?', (username, hash_password(password)))
    data = c.fetchone()
    return data

def hash_password(password):
    """密码哈希处理"""
    return hashlib.sha256(password.encode()).hexdigest()

def login_page():
    """登录页面"""
    st.title("登录")
    username = st.text_input("用户名", key="login_username")
    password = st.text_input("密码", type="password", key="login_password")
    if st.button("登录"):
        result = login_user(username, password)
        if result:
            st.session_state['logged_in'] = True
            st.session_state['username'] = username
            st.experimental_rerun()
        else:
            st.error("用户名或密码错误")

def register_page():
    """注册页面"""
    st.title("注册")
    username = st.text_input("用户名", key="register_username")
    password = st.text_input("密码", type="password", key="register_password")
    if st.button("注册"):
        if username and password:
            add_user(username, password)
            st.success("注册成功，请登录")
        else:
            st.error("用户名和密码不能为空")